{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Internal API\n",
    "This notebook is meant to provide insight into the inner workings of TrajectoryOptimization.jl. This is geared towards people who are interested in using TrajectoryOptimization to create a solver. It provides examples of how to use the methods provided in TrajectoryOptimization to evaluate the dynamics, objective, and constraints. It assumes a decent level of profficiency with Julia, a good knowledge of trajectory optimization, and some familiarity with the basics of TrajectoryOptimization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m\u001b[1m Activating\u001b[22m\u001b[39m environment at `~/.julia/dev/TrajectoryOptimization/examples/Project.toml`\n",
      "┌ Info: Precompiling TrajectoryOptimization [c79d492b-0548-5874-b488-5a62c1d9d0ca]\n",
      "└ @ Base loading.jl:1278\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "RobotDynamics"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import Pkg; Pkg.activate(@__DIR__)\n",
    "using TrajectoryOptimization\n",
    "using RobotDynamics\n",
    "using StaticArrays, LinearAlgebra\n",
    "using Rotations\n",
    "using RobotZoo: Quadrotor\n",
    "using Test\n",
    "using BenchmarkTools\n",
    "const TO = TrajectoryOptimization\n",
    "const RD = RobotDynamics"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create a Problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = Quadrotor()\n",
    "n,m = size(model)           # number of states and controls\n",
    "n̄ = state_diff_size(model)  # size of error state\n",
    "N = 51                      # number of knot points\n",
    "tf = 5.0                    # final time\n",
    "\n",
    "# initial and final conditions\n",
    "x0 = RBState([1,2,1], UnitQuaternion(I), zeros(3), zeros(3))\n",
    "xf = RBState([0,0,2], UnitQuaternion(I), zeros(3), zeros(3))\n",
    "\n",
    "# objective\n",
    "Q = Diagonal(@SVector fill(0.1, n))\n",
    "R = Diagonal(@SVector fill(0.01, m))\n",
    "Qf = Diagonal(@SVector fill(100.0, n))\n",
    "obj = LQRObjective(Q,R,Qf,xf,N)\n",
    "\n",
    "# constraints\n",
    "cons = ConstraintList(n,m,N)\n",
    "add_constraint!(cons, BoundConstraint(n,m, u_min=zeros(4), u_max=fill(10.0,4)), 1:N-1)\n",
    "add_constraint!(cons, CircleConstraint(n, SA_F64[1,2], SA_F64[1,2], SA[0.1,0.1]), 1:N-1)\n",
    "add_constraint!(cons, GoalConstraint(xf, SA[1,2,3]), N)\n",
    "\n",
    "# problem\n",
    "prob = Problem(model, obj, xf, tf, x0=x0, constraints=cons);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initialize the controls\n",
    "u0 = zeros(model)[2]                        # get hover control\n",
    "initial_controls!(prob, u0)                 # set all time-steps to the same\n",
    "initial_controls!(prob, [u0 for k = 1:N-1]) # use a vector of initial controls\n",
    "initial_controls!(prob, fill(u0[1],m,N-1))  # use a matrix of initial controls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulating the dynamics\n",
    "To simulate the dynamics forward from the initial state, use the `rollout!` method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13-element SArray{Tuple{13},Float64,1,13} with indices SOneTo(13):\n",
       " 1.0\n",
       " 2.0\n",
       " 1.5001876006215789\n",
       " 0.7209485318742003\n",
       " 0.0\n",
       " 0.0\n",
       " 0.6929869546740692\n",
       " 0.0\n",
       " 0.0\n",
       " 0.20003656032918066\n",
       " 0.0\n",
       " 0.0\n",
       " 0.6124999999999969"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# simulate the system forward\n",
    "rollout!(prob)\n",
    "@test state(prob.Z[1]) == prob.x0\n",
    "@test states(prob)[end] ≈ prob.x0\n",
    "\n",
    "# alternative method\n",
    "rollout!(TO.integration(prob), prob.model, prob.Z, prob.x0)\n",
    "@test state(prob.Z[1]) == prob.x0\n",
    "@test states(prob)[end] ≈ prob.x0\n",
    "\n",
    "# change control so that the state changes\n",
    "u0 += [1,0,1,0]*1e-2\n",
    "initial_controls!(prob, u0)\n",
    "rollout!(TO.integration(prob), prob.model, prob.Z, prob.x0)\n",
    "states(prob)[end]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Computing the dynamics Jacobians\n",
    "For normal models, all you just need to pre-allocate some `DynamicsExpansion` types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "D = [TO.DynamicsExpansion{Float64}(n,n̄,m) for k = 1:N-1]\n",
    "TO.dynamics_expansion!(TO.integration(prob), D, model, prob.Z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " For a `LieGroupModel`, we need to first compute the error-state Jacobians, and then use those along with the Jacobians we computed before to calculate the Jacobians on the error dynamics. Note that we allocate one extra for use as a temporary array for intermediate matrix multiplications."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "G = [SizedMatrix{n,n̄}(zeros(n,n̄)) for k = 1:N+1]\n",
    "RD.state_diff_jacobian!(G, model, prob.Z)\n",
    "TO.error_expansion!(D, model, G)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can extract the dynamics Jacobians as each time step using `TO.error_expansion`. Note that the size of these Jacobians match the size of the error state (12)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "12×12 SizedArray{Tuple{12,12},Float64,2,2,Array{Float64,2}} with indices SOneTo(12)×SOneTo(12):\n",
       " 1.0  0.0  0.0   0.0        0.0985     0.0  0.1  …   0.00164167   0.0\n",
       " 0.0  1.0  0.0  -0.0985     0.0        0.0  0.0      0.0          0.0\n",
       " 0.0  0.0  1.0   0.0        0.0        0.0  0.0      0.0          0.0\n",
       " 0.0  0.0  0.0   1.0        0.0006125  0.0  0.0      1.28714e-5   0.0\n",
       " 0.0  0.0  0.0  -0.0006125  1.0        0.0  0.0      0.05         0.0\n",
       " 0.0  0.0  0.0   0.0        0.0        1.0  0.0  …   0.0          0.05\n",
       " 0.0  0.0  0.0   0.0        1.97       0.0  1.0      0.04925      0.0\n",
       " 0.0  0.0  0.0  -1.97       0.0        0.0  0.0      5.0276e-6    0.0\n",
       " 0.0  0.0  0.0   0.0        0.0        0.0  0.0      0.0          5.0276e-6\n",
       " 0.0  0.0  0.0   0.0        0.0        0.0  0.0     -0.000452717  0.0\n",
       " 0.0  0.0  0.0   0.0        0.0        0.0  0.0  …   1.0          0.0\n",
       " 0.0  0.0  0.0   0.0        0.0        0.0  0.0      0.0          1.0"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "12×4 SizedArray{Tuple{12,4},Float64,2,2,Array{Float64,2}} with indices SOneTo(12)×SOneTo(4):\n",
       "  0.0          0.0          0.0          0.0\n",
       "  0.0          0.0          0.0          0.0\n",
       "  0.01         0.01         0.01         0.01\n",
       " -2.95492e-5   0.190217     2.95492e-5  -0.190217\n",
       " -0.190217    -2.95492e-5   0.190217     2.95492e-5\n",
       "  0.0153125   -0.0153125    0.0153125   -0.0153125\n",
       " -0.124909     7.6507e-5    0.124909    -7.6507e-5\n",
       " -7.6507e-5   -0.124909     7.6507e-5    0.124909\n",
       "  0.200006     0.199994     0.200006     0.199994\n",
       "  0.00229639   7.6087      -0.00229639  -7.6087\n",
       " -7.6087       0.00229639   7.6087      -0.00229639\n",
       "  0.6125      -0.6125       0.6125      -0.6125"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "A,B = TO.error_expansion(D[1], model)\n",
    "display(A)\n",
    "display(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing the Cost\n",
    "Getting the cost from the `Problem` type is very easy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "312.79752892757296"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cost(prob)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "However, under the hood there's a bit more going on. This method actually calls:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "TO.cost!(obj, prob.Z);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "which computes the cost for each time step, storing it in the vector `J` stored inside the objective. It can be helpful to have access to the vector when debugging the cost, so we provide a getter method:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×51 Adjoint{Float64,Array{Float64,1}}:\n",
       " 0.033032  0.0330308  0.0330273  0.0330215  …  0.0337857  0.0339594  311.154"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TO.get_J(obj)'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cost at a single knot point can be computed using `TO.stage_cost`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.033032003124999994"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TO.stage_cost(obj[1], prob.Z[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cost Expansion\n",
    "To store the 2nd order expansion of the cost function we use the `TO.Expansion` type, and the convenient `TO.CostExpansion` constructor:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array{TrajectoryOptimization.Expansion{13,4,Float64},1}"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "E0 = TO.CostExpansion(n, m, N)\n",
    "typeof(E0)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To evaluate the 2nd Order Taylor Series expansion of the cost we can use the following methods:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "true"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "TO.cost_expansion!(E0, obj, prob.Z)\n",
    "\n",
    "# this just calls these two functions\n",
    "TO.cost_gradient!(E0, obj, prob.Z)\n",
    "TO.cost_hessian!(E0, obj, prob.Z)\n",
    "\n",
    "# which call these functions each time step\n",
    "TO.gradient!(E0[1], obj[1], prob.Z[1])\n",
    "TO.hessian!(E0[1], obj[1], prob.Z[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The boolean flag returned from `TO.hessian!` specifies that the hessian is constant. Some computation can be saved by leveraging this information. Internally, this is stored in the objective:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×51 Adjoint{Bool,BitArray{1}}:\n",
       " 1  1  1  1  1  1  1  1  1  1  1  1  1  …  1  1  1  1  1  1  1  1  1  1  1  1"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "obj.const_hess'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can force it to recompute using the optional `init` flag:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  38.098 μs (0 allocations: 0 bytes)\n",
      "  2.591 μs (0 allocations: 0 bytes)\n"
     ]
    }
   ],
   "source": [
    "@btime TO.cost_expansion!($E0, $obj, $(prob.Z), init=true)\n",
    "@btime TO.cost_expansion!($E0, $obj, $(prob.Z), init=false)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `TO.cost_expansion!` takes one additional argument: `rezero`. If true, this will multiply each of the terms in the Hessian by zero prior to computing the new expansion. By default, the expansions may choose to only modify the elements they know to be non-zero (e.g. Diagonal costs only modify the diagonal). If we end up adding addition terms to the expansion (i.e. augmented Lagrangian penalty terms), our expansion will be incorrect if we don't re-zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(E0[1]).Q[1, 2] = 10.0\n",
      "(E0[1]).Q[1, 2] = 0.0\n"
     ]
    }
   ],
   "source": [
    "E0[1].Q[1,2] = 10\n",
    "TO.cost_expansion!(E0, obj, prob.Z)\n",
    "@show E0[1].Q[1,2]  # this is incorrect\n",
    "TO.cost_expansion!(E0, obj, prob.Z, init=true, rezero=true)\n",
    "@show E0[1].Q[1,2]; # now correct"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Expansion caches (advanced)\n",
    "TrajectoryOptimization provides methods to differentiate custom nonlinear cost functions automatically with either [ForwardDiff.jl](https://github.com/JuliaDiff/ForwardDiff.jl) or [FiniteDiff.jl](https://github.com/JuliaDiff/FiniteDiff.jl). For finite differencing, we can save compuation and allocations by passing a cache. This can be automatically allocated using `TO.ExpansionCache(cost::CostFunction)`. While the automatic methods (e.g. `_gradient!` and `_hessian!`) shouldn't be modified, the `TO.ExpansionCache` can be overwritten for a custom cost function type where you want to pre-allocated some memory to use while computing the gradient and Hessian of the cost function. \n",
    "\n",
    "This function returns a vector of 4 caches: `[grad, hess, grad_term, hess_term]` which are meant to be used by gradient and Hessian computations for the stage and terminal cost functions, respectively. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Error State Expansions\n",
    "If the model is a `LieGroupModel`, like our Quadrotor model, we need to again convert the expansions to be valid for the lower-dimensional error state. This ends up being very similar to what we did for the dynamics. First, we can use the following constructor to initialize a new `QuadraticObjective` to store the error state expansion. In the second line, this constructor will alias `E === E0` if model is not a `LieGroupModel`, saving computation and avoiding unnecessary storage."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "E = TO.CostExpansion(n̄, m, N)     # error state Jacobians\n",
    "E0 = TO.CostExpansion(E, model);  # \"normal\" Jacobians"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now compute the expansion on the error state using methods very similar to what we did previously for the dynamics Jacobian. Note that the 3x3 diagonal block for the rotation state is no longer diagonal."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×3 Array{Float64,2}:\n",
       " 0.00741626   2.65474e-19  0.0\n",
       " 2.84866e-19  0.00741626   0.0\n",
       " 0.0          0.0          0.00741626"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "RD.state_diff_jacobian!(G, model, prob.Z)\n",
    "TO.cost_expansion!(E0, obj, prob.Z, init=true, rezero=true)\n",
    "TO.error_expansion!(E, E0, model, prob.Z, G);\n",
    "E[N-1].Q[4:6,4:6]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Constraints\n",
    "Constraints in TrajectoryOptimization can get a little complicated. Let's start by making a couple important points:\n",
    "* Each constraint is defined by an `AbstractConstraint`. Specific per-timestep constraints should inherit from it's subtypes: `StageConstraint`, `ControlConstraint`, or `StageConstraint`.\n",
    "* Each constraint is associated with a set of time indices `inds`, which specifies the time steps for which the knotpoint applies.\n",
    "* `ConstraintList` only stores the constraint definition. It does not store constraint values and therefore cannot be evaluated. This is the representation stored in `Problem`.\n",
    "* `AbstractConstraintSet` is meant to instantiated by the solver, and will usually be a simple wrapper around a vector of `AbstractConstraintValues`.\n",
    "* The `AbstractConstraintValues` type is a wrapper around an `AbstractConstraint` and adds storage for constraint values, Jacobians, and anything else needed by the solver when dealing with constraints. This can implemented by the solver, or the solver can use the built-in `ConVal` type, demonstrated below.\n",
    "\n",
    "## Stage Constraints\n",
    "Let's start by looking at the `CircleConstraint` we created:\n",
    "```julia\n",
    "CircleConstraint(n, SA_F64[1,2], SA_F64[1,2], SA[0.1,0.1])\n",
    "```\n",
    "Note that this is a `StateConstraint`, meaning it is only a function of the state at a single time step. We also show a few methods for extracting information from the constraint."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "typeof(con) = CircleConstraint{2,Float64}\n",
      "Length of the constraint: 2\n"
     ]
    }
   ],
   "source": [
    "con,inds = cons[2], 1:N-1\n",
    "@show typeof(con)                     # just verifying it's a CircleConstraint\n",
    "@test con isa TO.StateConstraint      # it inherits from StateConstraint, so it's a function of a single state\n",
    "@test state_dim(con) == n             # the state dimension. control_dim won't be defined.\n",
    "@test TO.check_dims(con, n, m)        # useful method to check if a constraint is consistent with the problem sizes\n",
    "p = length(con);                      # get the length of the constraint vector\n",
    "println(\"Length of the constraint: $p\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This means that the `evaluate` can be passed a single state vector. Alternatively, we can pass a `KnotPoint` and the state will be extracted automatically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2-element SArray{Tuple{2},Float64,1,2} with indices SOneTo(2):\n",
       " -0.99\n",
       " -0.99"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z = prob.Z[end-1]\n",
    "x = state(z)\n",
    "v1 = TO.evaluate(con, x)\n",
    "v2 = TO.evaluate(con, z) \n",
    "@test v1 == v2\n",
    "v1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For the constraint Jacobian, since it's only a function of the state, we'd expect it to be of size `(p,n)`. A `StageConstraint` will have a size of `(p,n+m)`, and `ControlConstraint` will have a size of `(p,m)`. We can generate the Jacobian automatically using `TO.gen_jacobian`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×13 SizedArray{Tuple{2,13},Float64,2,2,Array{Float64,2}} with indices SOneTo(2)×SOneTo(13):\n",
       " 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0\n",
       " 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jac = TO.gen_jacobian(con)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now evaluate the constraint using `TO.jacobian!`. We can pass in either the state or a `KnotPoint`. As with cost Hessians, the constraint Jacobian returns a boolean specifying whether the Jacobian is constant or not (i.e. the constraint is linear)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "false"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "jacobian!(jac, con, x)\n",
    "jacobian!(jac, con, z)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dynamics Constraints\n",
    "TrajectoryOptimization also naturally has to deal with dynamics constraints, which couple states and controls across 2 adjacent time steps. Let's start by creating a dynamics constraint using the built-in `TO.add_dynamics_constraint!`. Internally, this is just creating a `DynamicsConstraint` and a `GoalConstraint` for the initial condition, calling the following constructor the `DynamicsConstraint`.\n",
    "```julia\n",
    "dyn_con = DynamicsConstraint{TO.integration(prob)}(prob.model, prob.N)\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of constraint before: 3\n",
      "Number of constraint after: 5\n",
      "Type of dynamics constraint:\n",
      "TrajectoryOptimization.DynamicsConstraint{RK3,Quadrotor{UnitQuaternion{Float64}},13,4,17,Float64}\n"
     ]
    }
   ],
   "source": [
    "prob2 = copy(prob)\n",
    "TO.add_dynamics_constraints!(prob2)\n",
    "println(\"Number of constraint before: \", length(TO.get_constraints(prob)))\n",
    "println(\"Number of constraint after: \", length(TO.get_constraints(prob2)))\n",
    "\n",
    "dyn = prob2.constraints[end]\n",
    "println(\"Type of dynamics constraint:\")\n",
    "println(typeof(dyn))  # type holds integration method, model type, and model sizes"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To evaluate the constraint, we need to pass in the states at the current and next time steps, along with the control at the current time step. We simplify this by simply requiring that we pass in the `KnotPoint`s for the current and next time step. As expected, the dynamics error for our rolled-out trajectory should be zero:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×13 Adjoint{Float64,SArray{Tuple{13},Float64,1,13}} with indices SOneTo(1)×SOneTo(13):\n",
       " 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "z2 = prob.Z[end]\n",
    "TO.evaluate(dyn, z, z2)'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The Jacobian gets a little trickier, since we need to separate the Jacobian between time steps. Our Jacobian looks something like:\n",
    "$$ [A \\; B \\;\\vert\\; -I] $$\n",
    "TrajectoryOptimization splits this by time step:\n",
    "$$ [\\nabla f_1 \\;\\vert\\; \\nabla f_2] $$\n",
    "\n",
    "To initialize the Jacobians, we can use the generic `TO.widths` which returns a tuple of \"widths,\" meaning the number of columns in each partition. Since our first partition has both state and controls, and the second only has states, we get:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(17, 13)"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "w = TO.widths(dyn)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can use this to initialize our Jacobians for both steps. When we call `jacobian!` we need to specify which of the time steps we're taking the Jacobian with respect to."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "MethodError: no method matching _discrete_jacobian!(::RobotDynamics.ForwardAD, ::Type{RK3}, ::Array{Float64,2}, ::Quadrotor{UnitQuaternion{Float64}}, ::RobotDynamics.GeneralKnotPoint{Float64,13,4,SArray{Tuple{17},Float64,1,17}}, ::FiniteDiff.JacobianCache{Array{Float64,1},Array{Float64,1},Array{Float64,1},UnitRange{Int64},Nothing,Val{:forward}(),Float64})\nClosest candidates are:\n  _discrete_jacobian!(::RobotDynamics.ForwardAD, ::Type{Q}, ::Any, ::AbstractModel, ::AbstractKnotPoint{T,N,M}) where {T, N, M, Q<:RobotDynamics.Explicit} at /home/brian/.julia/packages/RobotDynamics/yaan9/src/model.jl:246\n  _discrete_jacobian!(!Matched::RobotDynamics.FiniteDifference, ::Type{Q}, ::Any, ::AbstractModel, ::AbstractKnotPoint{T,N,M}, ::Any) where {T, N, M, Q<:RobotDynamics.Explicit} at /home/brian/.julia/packages/RobotDynamics/yaan9/src/model.jl:255\n  _discrete_jacobian!(!Matched::RobotDynamics.FiniteDifference, ::Type{Q}, ::Any, ::AbstractModel, ::AbstractKnotPoint{T,N,M}) where {T, N, M, Q<:RobotDynamics.Explicit} at /home/brian/.julia/packages/RobotDynamics/yaan9/src/model.jl:255",
     "output_type": "error",
     "traceback": [
      "MethodError: no method matching _discrete_jacobian!(::RobotDynamics.ForwardAD, ::Type{RK3}, ::Array{Float64,2}, ::Quadrotor{UnitQuaternion{Float64}}, ::RobotDynamics.GeneralKnotPoint{Float64,13,4,SArray{Tuple{17},Float64,1,17}}, ::FiniteDiff.JacobianCache{Array{Float64,1},Array{Float64,1},Array{Float64,1},UnitRange{Int64},Nothing,Val{:forward}(),Float64})\nClosest candidates are:\n  _discrete_jacobian!(::RobotDynamics.ForwardAD, ::Type{Q}, ::Any, ::AbstractModel, ::AbstractKnotPoint{T,N,M}) where {T, N, M, Q<:RobotDynamics.Explicit} at /home/brian/.julia/packages/RobotDynamics/yaan9/src/model.jl:246\n  _discrete_jacobian!(!Matched::RobotDynamics.FiniteDifference, ::Type{Q}, ::Any, ::AbstractModel, ::AbstractKnotPoint{T,N,M}, ::Any) where {T, N, M, Q<:RobotDynamics.Explicit} at /home/brian/.julia/packages/RobotDynamics/yaan9/src/model.jl:255\n  _discrete_jacobian!(!Matched::RobotDynamics.FiniteDifference, ::Type{Q}, ::Any, ::AbstractModel, ::AbstractKnotPoint{T,N,M}) where {T, N, M, Q<:RobotDynamics.Explicit} at /home/brian/.julia/packages/RobotDynamics/yaan9/src/model.jl:255",
      "",
      "Stacktrace:",
      " [1] discrete_jacobian!(::Type{RK3}, ::Array{Float64,2}, ::Quadrotor{UnitQuaternion{Float64}}, ::RobotDynamics.GeneralKnotPoint{Float64,13,4,SArray{Tuple{17},Float64,1,17}}, ::FiniteDiff.JacobianCache{Array{Float64,1},Array{Float64,1},Array{Float64,1},UnitRange{Int64},Nothing,Val{:forward}(),Float64}) at /home/brian/.julia/packages/RobotDynamics/yaan9/src/model.jl:243",
      " [2] jacobian!(::Array{Float64,2}, ::TrajectoryOptimization.DynamicsConstraint{RK3,Quadrotor{UnitQuaternion{Float64}},13,4,17,Float64}, ::RobotDynamics.GeneralKnotPoint{Float64,13,4,SArray{Tuple{17},Float64,1,17}}, ::RobotDynamics.GeneralKnotPoint{Float64,13,4,SArray{Tuple{17},Float64,1,17}}, ::Int64) at /home/brian/.julia/dev/TrajectoryOptimization/src/dynamics_constraints.jl:75",
      " [3] top-level scope at In[26]:3",
      " [4] include_string(::Function, ::Module, ::String, ::String) at ./loading.jl:1091"
     ]
    }
   ],
   "source": [
    "jac1 = zeros(n,w[1])  # can also use TO.gen_jacobian(dyn)\n",
    "jac2 = zeros(n,w[2])\n",
    "println(\"Is constant? \", jacobian!(jac1, dyn, z, z2, 1))\n",
    "println(\"Is constant? \", jacobian!(jac2, dyn, z, z2, 2))\n",
    "display(jac1)\n",
    "display(jac2);"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `AbstractConstraintValues` type\n",
    "The `AbstractConstraintValues` type is meant to store all of the constraint values and Jacobians. We'll use the provided `ConVal` instantiation here. This type allows us to change the underlying storage type of our constraint values and Jacobians, allowing us to either use small dense arrays, or views into a large, sparse matrix. We show how to do both below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "cons = prob2.constraints                      # get `ConstraintList` with dynamics constraint\n",
    "circle,inds = cons[3], cons.inds[3]\n",
    "dyn = cons[end]\n",
    "\n",
    "### Dense Matrix Storage ###\n",
    "\n",
    "# to allocate our storage, we can use the provided `TO.gen_convals` method\n",
    "p = length(circle)\n",
    "C,c = TO.gen_convals(n,m,circle,inds)\n",
    "@test length(C) == length(c) == length(inds)  # all these should be same length, the number of time steps\n",
    "@test size(C[1]) == (p,n)                     # each element in a Jacobian\n",
    "@test typeof(C[1]) <: SizedMatrix{p,n}        # each element in a SizedMatrix\n",
    "\n",
    "# we can now generate our ConVal type\n",
    "cval = TO.ConVal(n, m, circle, inds, C, c)\n",
    "\n",
    "\n",
    "### Sparse Matrix Storage ###\n",
    "using SparseArrays\n",
    "\n",
    "# allocate sparse array and vector\n",
    "NN = TO.num_vars(prob.Z)                      # total number of states and control (decision variables)\n",
    "P = sum(TO.num_constraints(cons))             # total number of constraints\n",
    "D = spzeros(P,NN)\n",
    "d = zeros(P)\n",
    "\n",
    "# to specify the ordering of the constraints, we use the provided `JacobianStructure` type\n",
    "jac_struct = TO.JacobianStructure(cons, :by_knotpoint)  # 2nd argument optional. Other option is :by_constraint\n",
    "\n",
    "# we can impart the structure to our sparse matrix using\n",
    "D_struct = spzeros(Int,P,NN)\n",
    "TO.jacobian_structure!(D_struct, jac_struct)\n",
    "\n",
    "# with our structure specified, we can now create views for all of our constraints at the same time\n",
    "C,c = TO.gen_convals(D,d,cons,jac_struct)\n",
    "\n",
    "# this output is a little different than the previous one, since it's for all of the constraints\n",
    "#   let's look at the output for the CircleConstraint (index 3)\n",
    "@test length(C) == length(cons)               # length of output is the number of constraints\n",
    "@test length(C[3]) == length(inds)            # length of each element is the number of time steps\n",
    "@test size(C[3][1]) == (p,n)                  # each sub-element is a Jacobian\n",
    "\n",
    "# for the DynamicsConstraint it's slightly different\n",
    "@test size(C[end]) == (50,2)                  # the Jacobians are stored in a Matrix\n",
    "@test size(c[end]) == (50,)                   # the constraint values are still stored in a Vector\n",
    "@test size(C[end][1,1]) == (n,n+m)            # the 1st column is the Jacobian for the current time step\n",
    "@test size(C[end][1,2]) == (n,n)              # the 2nd column is the Jacobian for the next time step\n",
    "@test typeof(C[end][1]) <: SubArray           # each Jacobian is a view into D\n",
    "@test typeof(c[end][1]) <: SubArray           # each constraint value is a view into d\n",
    "\n",
    "# let's verify the views worked\n",
    "c[end][1] .= 1\n",
    "@test d[jac_struct.cinds[end][1]] == ones(n)  # the jac_struct stores the indices we need\n",
    "C[end][2,1] .= 2\n",
    "@test D[jac_struct.cinds[end][2], jac_struct.zinds[end][2]] == fill(2,n,n+m)\n",
    "C[end][2,2] .= 3\n",
    "@test D[jac_struct.cinds[end][2], jac_struct.zinds[end][3]] == [fill(3,n,n) zeros(n,m)]\n",
    "\n",
    "# we can now generate a list of ConVals\n",
    "cvals = map(1:length(cons)) do i\n",
    "    TO.ConVal(n,m,cons[i],cons.inds[i],C[i],c[i])\n",
    "end;"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "With an `AbstractContraintValues` type, we're now provided a few convenient methods for evaluating our constraint:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2×13 SizedArray{Tuple{2,13},Float64,2,2,Array{Float64,2}} with indices SOneTo(2)×SOneTo(13):\n",
       " -0.0  -2.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0\n",
       "  2.0  -0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Calculate the constraint values\n",
    "TO.evaluate!(cval, prob.Z)\n",
    "\n",
    "# Calculate the constraint Jacobian\n",
    "TO.jacobian!(cval, prob.Z)\n",
    "\n",
    "# Calcualate the maximum constraint violation\n",
    "max_violation(cval)\n",
    "\n",
    "# Individual constraint values and Jacobians are accessed via the `vals` and `jac` fields\n",
    "cval.vals[2]   # constraint value at 2nd time step\n",
    "cval.jac[2,1]  # constraint Jacobian at the 2nd time step, for the current time step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is easier than calling the underlying methods that work across all time steps. Note that these methods can be helpful to know about, since they can be overloaded for a custom constraint that needs to change behavior at different time steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [],
   "source": [
    "TO.evaluate!(cval.vals, circle, prob.Z, inds)\n",
    "TO.jacobian!(cval.jac, circle, prob.Z, inds, cval.is_const)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `AbstractConstraintSet`\n",
    "Now that we're familiar with the `AbstractConstraintValues` type, we can basically put these together into a vector and call it an `AbstractConstraintSet`. We'll create a minimal constraint set here.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct BasicConstraintSet{T} <: TO.AbstractConstraintSet\n",
    "    convals::Vector{TO.ConVal}\n",
    "    c_max::Vector{T}\n",
    "end\n",
    "@inline TO.get_convals(conSet::BasicConstraintSet) = conSet.convals\n",
    "\n",
    "conSet = BasicConstraintSet(Vector{TO.ConVal}(cvals), zeros(length(cvals)));"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.0"
      ]
     },
     "execution_count": 39,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "@test length(conSet) == length(cons)  # supports iteration (currently indexing is not supported)\n",
    "\n",
    "# supports the same methods we used on `ConVal`, but just applies them to all constraints\n",
    "TO.evaluate!(conSet, prob.Z)\n",
    "TO.jacobian!(conSet, prob.Z)\n",
    "max_violation(conSet)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.5.3",
   "language": "julia",
   "name": "julia-1.5"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
